const tf = require('@tensorflow/tfjs-node');
const getData = require('./load');

const TRAIN_DIR = '垃圾分类/train'
const OUTPUT_DIR = 'output';
const MOBILENET_URL = 'http://ai-sample.oss-cn-hangzhou.aliyuncs.com/pipcook/models/mobilenet/web_model/model.json';
const main = async () => {
    console.log('It is main');
    // 1、加载数据
    const { xs, ys, classes } = await getData(TRAIN_DIR, OUTPUT_DIR);
    console.log(xs, ys, classes);

    // 2、定义模型
    const mobilenet = await tf.loadLayersModel(MOBILENET_URL);
    mobilenet.summary();
    // 截断模型
    const model = tf.sequential();
    for (let i = 0; i <= 85; i++) {
        const layer = mobilenet.layers[i];
        layer.trainable = false;
        model.add(layer);
    }
    // console.log(mobilenet.layers.map((layer, i) => [layer.name, i]));
    // 数据维度摊平
    model.add(tf.layers.flatten());
    // 增加隐藏层
    model.add(tf.layers.dense({
        units: 10,
        activation: 'relu',
    }));
    // 增加输出层
    model.add(tf.layers.dense({
        units: classes.length,
        activation: 'softmax',
    }));


    // 3、训练模型
    model.compile({
        loss: 'sparseCategoricalCrossentropy', // 损失函数
        optimizer: tf.train.adam(), // 优化器
        metrics: ['acc'] // 精确度
    });
    await model.fit(xs, ys, { epochs: 20 });
    await model.save(`file://${process.cwd()}/${OUTPUT_DIR}`);
}

main();